
package org.knot.mockejb3;

import java.security.Identity;
import java.util.Properties;

import javax.ejb.EJBException;
import javax.ejb.EJBHome;
import javax.ejb.EJBLocalHome;
import javax.ejb.EJBLocalObject;
import javax.ejb.EJBObject;
import javax.transaction.Status;
import javax.transaction.SystemException;
import javax.transaction.UserTransaction;

import org.knot.mockejb3.exception.MethodNotImplementedException;
import org.knot.mockejb3.transcation.TransactionManager;


/**
 * Provides implementation of <code>javax.ejb.SessionContext</code>,
 * <code>javax.ejb.MessageDrivenContext</code> and <code>javax.ejb.EntityContext</code>
 * as well as some extra convenience methods.
 * 
 */
public class MockEjbContext {

    private Object homeProxy;
    private Object ejbObjectProxy;
    private Object primaryKey;
    
    
    MockEjbContext( final Object homeProxy ){
        this.homeProxy = homeProxy;
    }


    void setEjbObjectProxy( final Object ejbObjectProxy ){
        this.ejbObjectProxy = ejbObjectProxy;
    }

    
    /**
     * Always returns empty Properties object. Bean-scoped environment 
     * is not supported directly.
     * @see javax.ejb.EJBContext#getEnvironment()
     */
    public Properties getEnvironment() {
        
        return new Properties();
    }


    /**
     * Calls {@link TransactionManager} to get the
     * <code>javax.transaction.UserTransaction</code> object.
     * @return <code>javax.transaction.UserTransaction</code> object  
     * @see javax.ejb.EJBContext#getUserTransaction()
     */
    public UserTransaction getUserTransaction() throws IllegalStateException {

        UserTransaction tran = TransactionManager.getUserTransaction();
        
        return tran;
    }

    /**
     * @see javax.ejb.EJBContext#setRollbackOnly()
     */
    public void setRollbackOnly() throws IllegalStateException {
        UserTransaction tran = getUserTransaction();
        try {
            tran.setRollbackOnly();
        }
        catch( SystemException se ){
            throw new EJBException("Error trying to call setRollbackOnly on transaction", se); 
        }

    }

    /**
     * @see javax.ejb.EJBContext#getRollbackOnly()
     */
    public boolean getRollbackOnly() throws IllegalStateException {
        
        int status = Status.STATUS_UNKNOWN;
        UserTransaction tran = getUserTransaction();
        try {
            status = tran.getStatus();   
        }
        catch( SystemException se ){
            throw new EJBException("Error trying to call getStatus on transaction", se); 
        }
        
        return( status == Status.STATUS_MARKED_ROLLBACK );
    }

    /**
     * Obtains a reference to the EJB local object that is currently associated with the instance.
     * @return the EJB object currently associated with the instance
     */
    public EJBLocalObject getEJBLocalObject() throws IllegalStateException {
        return (EJBLocalObject)ejbObjectProxy;
    }

    /**
     * Obtains a reference to the EJB object that is currently associated with the instance.
     * @return the EJB object currently associated with the instance
     */
    public EJBObject getEJBObject() throws IllegalStateException {
        return (EJBObject)ejbObjectProxy;
    }
    
    // EJB 2.1 methods for the future 
    /*
    public javax.xml.rpc.handler.MessageContext getMessageContext() throws IllegalStateException{
        throwMethodNotImplemented( "getMessageContext");
        return null;
    }
    
    public TimerService getTimerService() throws IllegalStateException {    
        throwMethodNotImplemented( "getMessageContext");
        return null;
    }
    */
    


	/**
     * Returns the primary key for entity beans. If the context is assotiated with 
     * Session bean or MDB, returns null. 
     * MockEJB does not automatically handles primary keys for 
     * entity beans. Since "create" returns null, you need to intercept "create" methods
     * of your entity bean and return the real PK. 
     * 
	 * @see javax.ejb.EntityContext#getPrimaryKey()
	 */
	public Object getPrimaryKey() throws IllegalStateException {
        return primaryKey;
	}
    
    public void setPrimaryKey( Object primaryKey ){
    	this.primaryKey = primaryKey;
    }



}
